#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <jsoncpp/json/json.h>

#include <thread>
#include <chrono>
#include <mutex>
#include <atomic>
#include <iostream>
#include <iomanip>

#define SERVER_IP   "127.0.0.1"
#define SERVER_PORT 12345

// --- 消息类型定义 ---
#define TEST 1000 // 对应服务器端的TEST消息类型
#define MSG_PING 9998 // 心跳请求
#define MSG_PONG 9999 // 心跳响应

// --- 心跳配置 ---
const int HEARTBEAT_INTERVAL_S = 15; // 15秒发送一次心跳

// --- 全局互斥锁 ---
std::mutex g_socket_mutex;  // 保护 socket 的写操作
std::mutex g_log_mutex;     // 保护日志输出

// --- 日志输出函数 ---
void log_print(const std::string& thread_name, const std::string& level, const std::string& message) {
    std::lock_guard<std::mutex> lock(g_log_mutex);
    
    // 获取当前时间
    auto now = std::chrono::system_clock::now();
    auto time_t = std::chrono::system_clock::to_time_t(now);
    auto tm = *std::localtime(&time_t);
    
    // 获取毫秒数
    auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
        now.time_since_epoch()) % 1000;
    
    // 格式化输出: [时间] [线程] [级别] 消息
    std::cout << "[" << std::put_time(&tm, "%H:%M:%S") 
              << "." << std::setfill('0') << std::setw(3) << ms.count() << "] "
              << "[" << std::setw(8) << std::left << thread_name << "] "
              << "[" << std::setw(4) << level << "] " 
              << message << std::endl;
}

// 便捷的日志宏
#define LOG_INFO(thread, msg) log_print(thread, "INFO", msg)
#define LOG_WARN(thread, msg) log_print(thread, "WARN", msg)
#define LOG_ERROR(thread, msg) log_print(thread, "ERROR", msg)

// 从文件描述符 'fd' 上完整地读取 'count' 字节到 'buf'
static ssize_t read_full(int fd, void *buf, size_t count) {
    size_t  off = 0;
    ssize_t n;
    while (off < count) {
        n = read(fd, (char*)buf + off, count - off);
        if (n < 0) {
            if (errno == EINTR) continue;
            return -1;
        }
        if (n == 0) {
            return off;
        }
        off += n;
    }
    return off;
}

// --- 消息发送函数 ---
bool send_message(int sock, uint16_t msg_type, const std::string& data) {
    size_t body_len = data.length();
    size_t total_len = 4 + body_len;

    if (body_len > 2040) { 
        LOG_ERROR("SEND", "消息过长，最大 2040 字节");
        return false;
    }

    uint8_t *buf = new uint8_t[total_len];
    if (!buf) {
        LOG_ERROR("SEND", "内存分配失败");
        return false;
    }

    uint16_t net_type = htons(msg_type);
    uint16_t net_len  = htons(body_len);

    memcpy(buf, &net_type, sizeof(net_type));
    memcpy(buf + 2, &net_len,  sizeof(net_len));
    if (body_len > 0) {
        memcpy(buf + 4, data.c_str(), body_len);
    }
    
    std::lock_guard<std::mutex> lock(g_socket_mutex);
    if (write(sock, buf, total_len) != (ssize_t)total_len) {
        LOG_ERROR("SEND", "发送数据失败: " + std::string(strerror(errno)));
        delete[] buf;
        return false;
    }
    delete[] buf;
    return true;
}

// --- 心跳线程执行函数 ---
void heartbeat_thread_func(int sock, std::atomic<bool>& keep_running) {
    LOG_INFO("HEART", "心跳线程启动，间隔 " + std::to_string(HEARTBEAT_INTERVAL_S) + " 秒");
    
    while (keep_running) {
        std::this_thread::sleep_for(std::chrono::seconds(HEARTBEAT_INTERVAL_S));

        if (!keep_running) break;

        Json::Value root;
        root["type"] = "MSG_PING";
        root["content"] = "ping from client";
        root["timestamp"] = static_cast<Json::Int64>(time(nullptr));

        std::string json_data = root.toStyledString();
        // 移除末尾的换行符，让日志更简洁
        if (!json_data.empty() && json_data.back() == '\n') {
            json_data.pop_back();
        }

        LOG_INFO("HEART", "发送心跳 PING");
        if (!send_message(sock, MSG_PING, json_data)) {
            LOG_ERROR("HEART", "发送心跳失败，线程退出");
            break; 
        }
    }
    LOG_INFO("HEART", "心跳线程已停止");
}

// --- 读取线程执行函数 ---
void read_thread_func(int sock, std::atomic<bool>& keep_running) {
    LOG_INFO("READ", "读取线程启动，等待服务器消息");
    
    while (keep_running) {
        uint16_t resp_type, resp_len;
        ssize_t n;
        
        n = read_full(sock, &resp_type, sizeof(resp_type));
        if (n <= 0) {
            if (n < 0 && keep_running) {
                LOG_ERROR("READ", "读取消息类型失败: " + std::string(strerror(errno)));
            }
            LOG_WARN("READ", "服务器连接断开");
            keep_running = false;
            break;
        }
        
        n = read_full(sock, &resp_len, sizeof(resp_len));
        if (n <= 0) {
            if (n < 0 && keep_running) {
                LOG_ERROR("READ", "读取消息长度失败: " + std::string(strerror(errno)));
            }
            LOG_WARN("READ", "服务器连接断开");
            keep_running = false;
            break;
        }
        
        resp_type = ntohs(resp_type);
        resp_len = ntohs(resp_len);
        
        // 处理 PONG 响应
        if (resp_type == MSG_PONG) {
            LOG_INFO("READ", "收到服务器 PONG 响应");
            continue;
        }

        // 接收其他消息类型的响应体
        uint8_t *body = (uint8_t *)malloc(resp_len + 1);
        if (!body) {
            LOG_ERROR("READ", "分配内存失败");
            break;
        }
        
        if (read_full(sock, body, resp_len) != resp_len) {
            LOG_ERROR("READ", "读取消息体失败");
            free(body);
            break;
        }
        body[resp_len] = '\0';

        // 解析并显示 JSON 响应
        Json::Value resp_root;
        Json::Reader reader;
        if (reader.parse((char*)body, resp_root)) {
            std::string msg = "收到服务器响应 [ID:" + std::to_string(resp_root["id"].asInt()) + 
                             "] 内容: " + resp_root["data"].asString();
            LOG_INFO("READ", msg);
        } else {
            std::string msg = "收到服务器消息 [类型:" + std::to_string(resp_type) + 
                             "] 长度:" + std::to_string(resp_len) + 
                             " 内容: " + std::string((char*)body);
            LOG_INFO("READ", msg);
        }

        free(body);
    }
    LOG_INFO("READ", "读取线程已停止");
}

int main() {
    int sock;
    struct sockaddr_in srv_addr;

    LOG_INFO("MAIN", "客户端程序启动");

    if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        LOG_ERROR("MAIN", "创建套接字失败: " + std::string(strerror(errno)));
        exit(EXIT_FAILURE);
    }

    memset(&srv_addr, 0, sizeof(srv_addr));
    srv_addr.sin_family = AF_INET;
    srv_addr.sin_port = htons(SERVER_PORT);
    if (inet_pton(AF_INET, SERVER_IP, &srv_addr.sin_addr) <= 0) {
        LOG_ERROR("MAIN", "IP地址转换失败");
        close(sock);
        exit(EXIT_FAILURE);
    }

    if (connect(sock, (struct sockaddr*)&srv_addr, sizeof(srv_addr)) < 0) {
        LOG_ERROR("MAIN", "连接服务器失败: " + std::string(strerror(errno)));
        close(sock);
        exit(EXIT_FAILURE);
    }
    LOG_INFO("MAIN", "成功连接到服务器 " + std::string(SERVER_IP) + ":" + std::to_string(SERVER_PORT));

    // --- 启动线程 ---
    std::atomic<bool> keep_running(true);
    std::thread hb_thread(heartbeat_thread_func, sock, std::ref(keep_running));
    std::thread read_thread(read_thread_func, sock, std::ref(keep_running));

    LOG_INFO("MAIN", "所有线程已启动，进入交互模式");
    LOG_INFO("MAIN", "请输入消息内容（直接回车退出）");

    // 主循环用于发送消息
    while (keep_running) {
        char *line = NULL;
        size_t len = 0;
        
        // 显示提示符（不使用日志系统，避免格式干扰）
        {
            std::lock_guard<std::mutex> lock(g_log_mutex);
            std::cout << "\n>>> ";
            std::cout.flush();
        }

        if (getline(&line, &len, stdin) < 0) {
            free(line);
            break; 
        }

        size_t input_len = strcspn(line, "\r\n");
        if (input_len == 0) {
            // 用户输入空行，退出
            free(line);
            LOG_INFO("MAIN", "用户请求退出");
            break;
        }
        
        line[input_len] = '\0';
        
        Json::Value root;
        root["id"] = TEST;
        root["data"] = std::string(line, input_len);
        
        Json::FastWriter writer;
        std::string json_str = writer.write(root);
        
        if (!json_str.empty() && json_str.back() == '\n') {
            json_str.pop_back();
        }
        
        LOG_INFO("MAIN", "发送消息: " + std::string(line, input_len));

        if (!send_message(sock, TEST, json_str)) {
            LOG_ERROR("MAIN", "发送消息失败");
            free(line);
            break;
        }
        
        free(line);
        line = NULL;
    }
    
    // --- 清理和退出 ---
    LOG_INFO("MAIN", "正在关闭连接并停止线程");
    keep_running = false;
    
    close(sock);

    if (read_thread.joinable()) {
        read_thread.join();
    }
    if (hb_thread.joinable()) {
        hb_thread.join();
    }

    LOG_INFO("MAIN", "程序已正常退出");
    return 0;
}